import os

from torch import nn, optim
import torch
from torch.utils.data import DataLoader
from dataset import *
from yolo_v3_net import *
from torch.utils.tensorboard import SummaryWriter


def loss_fun(output, target, c):
    # target (13, 13, 3, 8) 其中，8这个维度，第一个是置信度，后面是4个坐标值，最后是三分类结果
    # 原output [1, 45, 13, 13] 修改为 N H W 3 15  因为我们需要取后面的值，这样定义方便取后面的值
    # permute() 是张量对象的方法，用于重新排列张量的维度。
    # 通过重新排列维度，可以改变张量的形状和维度的顺序。这通常用于满足特定计算需求或与其他张量进行操作时的匹配。
    output = output.permute(0, 2, 3, 1)   # [1, 13, 13, 45]
    # [1, 13, 13, 3, 15]
    output = output.reshape(output.size(0), output.size(1), output.size(2), 3, -1)
    # target[..., 0] 是一个索引操作符，用于从 target 张量中选择特定的元素。
    # 切片操作： ... 可以用作切片操作符的一部分，用于表示多个维度的连续切片。
    # 例如，my_array[..., 1:4] 表示对 my_array 数组的所有维度进行切片，但只在最后一个维度上选取索引从 1 到 3（不包含）的子集。
    # [..., 0] 表示我们要选择 target 张量中的所有元素在最后一个维度上的索引为 0 的子集。我们从 target 张量中选择了所有元素的最后一个维度中索引为 0 的元素。
    # target[..., 0] > 0 是一个比较操作，用于检查选择的子集是否大于 0。这将返回一个布尔张量，其中元素的值为 True 或 False
    # 因此，mask_obj 是一个布尔张量，其中的值表示 target 张量中最后一个维度索引为 0 的元素是否大于 0。

    # 正样本
    mask_obj = target[..., 0] > 0
    # 负样本
    mask_no_obj = target[..., 0] == 0

    # 置信度损失
    # loss_p_fun 是一个使用二进制交叉熵损失函数（BCELoss）初始化的损失函数对象。
    loss_p_fun = nn.BCELoss()
    # 应用 sigmoid 函数。这是为了将输出张量的值压缩到 [0, 1] 的范围内，以便满足二进制交叉熵损失函数的输入要求。
    # 这样不仅方便计算，而且因为模型最开始训练时，置信度都是随机值。将其值约束在0-1之间，不会出现计算精度的问题，梯度爆炸的情况
    loss_p = loss_p_fun(torch.sigmoid(output[..., 0]), target[..., 0])

    # bbox回归损失，应用于正样本，所以我们要取包含正样本目标的数据
    loss_box_fun = nn.MSELoss()
    loss_box = loss_box_fun(output[mask_obj][..., 1:5], target[mask_obj][..., 1:5])

    # 多分类损失
    # nn.CrossEntropyLoss() 是一个 PyTorch 提供的用于计算多类别分类任务的损失函数。它将输入的预测和目标进行比较，并计算交叉熵损失。
    # 在使用 loss_segment_fun 进行损失计算时，通常需要注意以下几点：
    # 输入的预测值应该是未经过 softmax 或 log softmax 处理的原始输出。交叉熵损失函数会在内部对预测值进行 softmax 处理。
    # 目标值（即真实标签）应该是一个以整数索引表示的张量。
    loss_segment_fun = nn.CrossEntropyLoss()
    loss_segment = loss_segment_fun(output[mask_obj][..., 5:],
                                    torch.argmax(target[mask_obj][..., 5:], dim=1, keepdim=True).squeeze(dim=1))

    # 在制作数据集时，正负样本的数量不均衡，一般均衡的情况为1比2，现在可能1比20几。对这种问题进行处理，在训练置信度这块，给它一个更大的损失权重。
    loss = c * loss_p + (1 - c) * 0.5 * loss_box + (1 - c) * 0.5 * loss_segment
    return loss

if __name__ == '__main__':
    summary_writer=SummaryWriter('logs')
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    dataset = YoloDataSet()
    # DataLoader 是一个用于加载数据的类，用于将数据集划分为小批量的样本进行训练。
    # 在每一轮训练中，DataLoader 会从数据集中随机选择 8 个样本并作为一个批次进行训练。
    # shuffle=True 表示在每个 epoch（训练周期）开始时，DataLoader 会将数据集中的样本打乱顺序，以增加数据的随机性和泛化能力。
    data_loader = DataLoader(dataset, batch_size=4, shuffle=True)

    weight_path = 'params/net34.pt'
    net = Yolo_V3_Net().to(device)
    if os.path.exists(weight_path):
        # torch.load(weight_path) 会从指定的 weight_path 路径加载保存的参数状态字典（state dictionary）。这个字典通常包含了模型的权重和偏置等参数。
        # net.load_state_dict() 方法将加载的参数状态字典加载到模型中，用于初始化模型的参数。
        net.load_state_dict(torch.load(weight_path))

    opt = optim.Adam(net.parameters()) # net.parameters() 是一个方法，用于获取模型中需要训练的参数列表。

    epoch = 1
    while True:
        index = 0
        total_loss = 0.0;
        for target_13, target_26, target_52, img_data in data_loader:
            target_13, target_26, target_52, img_data = target_13.to(device), target_26.to(device), target_52.to(
                device), img_data.to(device)

            output_13, output_26, output_52 = net(img_data)

            loss_13 = loss_fun(output_13.float(), target_13.float(), 0.5)
            loss_26 = loss_fun(output_26.float(), target_26.float(), 0.5)
            loss_52 = loss_fun(output_52.float(), target_52.float(), 0.5)

            loss=loss_13+loss_26+loss_52

            # zero_grad() 方法用于将优化器中所有参数的梯度（gradient）置零。在进行反向传播和梯度更新之前，我们需要清零之前计算的梯度，以避免梯度的累加影响。
            opt.zero_grad()
            # backward() 方法用于执行反向传播算法，计算相对于张量中各个参数的梯度（gradient）。在反向传播过程中，梯度会从损失张量向后传播，通过链式法则计算相对于每个参数的梯度。
            loss.backward()
            # 通过调用 optimizer.step() 方法，可以根据计算的梯度值对这些参数进行更新。
            opt.step()

            total_loss += loss.item();
            index+=1
        print(f'loss{epoch}==', total_loss)
        # summary_writer.add_scalar('train_loss', loss, index)

        # torch.save() 是 PyTorch 提供的函数，用于将模型的状态字典（state_dict）保存到文件中。
        # net.state_dict() 是模型对象 net 的方法，用于返回一个包含了模型所有参数的字典。模型的状态字典中包含了模型中所有需要学习的参数的值。
        # 'params/net.pt' 是保存模型状态字典的文件路径。
        my_path = 'params/net' + str(epoch) + '.pt'
        torch.save(net.state_dict(), my_path)
        print('模型保存成功')
        epoch+=1